package cirno.open.api.dao;

import cirno.open.api.config.DockerConfig;
import cirno.open.api.entity.ContainerInfo;
import cirno.open.api.entity.ExecuteResponse;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.CreateContainerCmd;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.api.model.HostConfig;
import com.github.dockerjava.api.model.StreamType;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * docker 操作类
 *
 * @author FlanChan
 * @since 2024/11/6
 */
@Slf4j
@Data
@Component
@RequiredArgsConstructor
public class DockerDao {

    private final DockerConfig dockerConfig;

    private final DockerClient dockerClient;

    /**
     * 执行命令
     *
     * @param containerId 容器 ID
     * @param cmd         CMD
     * @return {@link ExecuteResponse}
     */
    public ExecuteResponse execCmd(String containerId, String[] cmd) {
        // 正常返回信息
        ByteArrayOutputStream resultStream = new ByteArrayOutputStream();
        // 错误信息
        ByteArrayOutputStream errorResultStream = new ByteArrayOutputStream();

        // 结果
        final boolean[] result = {true};
        final boolean[] timeout = {true};
        try (ResultCallback.Adapter<Frame> frameAdapter = new ResultCallback.Adapter<Frame>() {

            @Override
            public void onComplete() {
                // 是否超时
                timeout[0] = false;
                super.onComplete();
            }

            @Override
            public void onNext(Frame frame) {
                StreamType streamType = frame.getStreamType();
                byte[] payload = frame.getPayload();
                if (StreamType.STDERR.equals(streamType)) {
                    try {
                        result[0] = false;
                        errorResultStream.write(payload);
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                } else {
                    try {
                        resultStream.write(payload);
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
                super.onNext(frame);
            }
        }) {



            ExecCreateCmdResponse execCompileCmdResponse = dockerClient.execCreateCmd(containerId)
                    .withCmd(cmd)
                    .withAttachStderr(true)
                    .withAttachStdin(true)
                    .withAttachStdout(true)
                    .exec();
            String execId = execCompileCmdResponse.getId();
            dockerClient.execStartCmd(execId).exec(frameAdapter).awaitCompletion(dockerConfig.getTimeoutLimit(), dockerConfig.getTimeUnit());

            // 超时
            if (timeout[0]) {
                return ExecuteResponse
                        .builder()
                        .success(false)
                        .errorMessage("执行超时")
                        .build();
            }

            return ExecuteResponse
                    .builder()
                    .success(result[0])
                    .message(resultStream.toString())
                    .errorMessage(errorResultStream.toString())
                    .build();

        } catch (IOException | InterruptedException e) {
            return ExecuteResponse
                    .builder()
                    .success(false)
                    .errorMessage(e.getMessage())
                    .build();
        }
    }

    public ContainerInfo startContainer(String codePath) {
        CreateContainerCmd containerCmd = dockerClient.createContainerCmd(dockerConfig.getImage());
        HostConfig hostConfig = new HostConfig();
        hostConfig.withMemory(dockerConfig.getMemoryLimit());
        hostConfig.withMemorySwap(dockerConfig.getMemorySwap());
        hostConfig.withCpuCount(dockerConfig.getCpuCount());
//        hostConfig.withReadonlyRootfs(true);
//        hostConfig.setBinds(new Bind(codePath, new Volume("/box")));
        CreateContainerResponse createContainerResponse = containerCmd
                .withHostConfig(hostConfig)
                .withNetworkDisabled(true)
                .withAttachStdin(true)
                .withAttachStderr(true)
                .withAttachStdout(true)
                .withTty(true)
                .exec();
        String containerId = createContainerResponse.getId();
        log.info("containerId: {}", containerId);
        // 启动容器
        dockerClient.startContainerCmd(containerId).exec();
        return ContainerInfo
                .builder()
                .containerId(containerId)
                .codePathName(codePath)
                .lastActivityTime(System.currentTimeMillis())
                .build();
    }

    /**
     * 复制文件到容器
     *
     * @param codeFile    代码文件
     * @param containerId 容器 ID
     */
    public void copyFileToContainer(String containerId, String codeFile) {
        dockerClient.copyArchiveToContainerCmd(containerId)
                .withHostResource(codeFile)
                .withRemotePath("/box")
                .exec();
    }

    public void cleanContainer(String containerId) {
        dockerClient.stopContainerCmd(containerId).exec();
        dockerClient.removeContainerCmd(containerId).exec();
    }
}
